#!/usr/bin/env python
# coding=utf-8

from argparse import ArgumentParser
import torch
import numpy as np
from pytorch_lightning import LightningModule, Trainer
from pytorch_lightning.callbacks import ModelCheckpoint

from models import Model

def cli_main():
	# args
	parser = ArgumentParser()
	parser = Trainer.add_argparse_args(parser)
	parser = Model.add_model_specific_args(parser)
	args = parser.parse_args()

	# get number of speakers (classes)
	utt2spk = np.loadtxt(args.train_utt2spk, dtype=np.str)
	spk_set = set()
	for utt, spk in utt2spk:
		spk_set.add(spk)
	args.num_classes = len(spk_set)

	# model
	model = Model(**vars(args))

	# config
	checkpoint_callback = ModelCheckpoint(monitor='val_loss',save_top_k=args.save_top_k, filename="{epoch}_{val_loss:.2f}")
	args.checkpoint_callback = checkpoint_callback

	# training
	trainer = Trainer.from_argparse_args(args)
	trainer.fit(model)

if __name__ == '__main__':  # pragma: no cover
	cli_main()

